#include "videodecode.h"
#include "common/nalu.h"
#include <math.h>

#ifdef __cplusplus
extern "C" {
#endif
    
static void log_callback_report(void *ptr, int level, const char *fmt, va_list vl)
{
    va_list vl2;
    char line[1024];
    static int print_prefix = 1;
    
    va_copy(vl2, vl);
    av_log_default_callback(ptr, level, fmt, vl);
    av_log_format_line(ptr, level, fmt, vl2, line, sizeof(line), &print_prefix);
    va_end(vl2);
    
    if (!strstr(line, "[h264"))
        LOGI("%s",line);
}


CVideoDecoder::CVideoDecoder(frame_available_cb on_frame_cb, error_cb on_error_cb, video_ready_cb on_ready_cb, statis_info_cb on_statis_cb)
: mOnFrameAvailableHandler(on_frame_cb), mOnErrorHandler(on_error_cb),
mOnVideoReadyHandler(on_ready_cb), mOnStatisInfoHandler(on_statis_cb)
{
    mState = STAT_IDLE;
    mDecodeCtx = NULL;
#ifdef ANDROID
    mSwsCtx = NULL;
#endif
    mUseExtraDecoder = false;
    mOnDecoderCreateHandler = NULL;
    
    mPacketCount = 0;
    mPacketWindex = 0;
    mPacketRindex = 0;
    memset(mPackets, 0, sizeof(mPackets));
    
    mFrameCount = 0;
    mFrameWindex = 0;
    mFrameRindex = 0;
    memset(mFrames, 0, sizeof(mFrames));
    
    memset(mExtradata, 0, sizeof(mExtradata));
    mExtradataSize = 0;
    
    mRtmpReader = new CRtmpReader();
    mRtmpReader->setPacketHandler(packet_handler, this);
    
    // 创建线程
    pthread_attr_t attr;
    pthread_attr_init(&attr);
    pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED);
    pthread_create(&mThread, &attr, thread_func, this);
    
    pthread_mutex_init(&mLock, 0);
}

CVideoDecoder::~CVideoDecoder()
{
    
    LOGE("~CVideoDecoder free ..");
    mIsRunning = false;
    if (mThread != 0) {
        for (int i = 0; i < 10; i++) {
            if (mThread == 0)
                break;
            usleep(1000 * 50);
        }
    }
    
    pthread_mutex_destroy(&mLock);
}

void CVideoDecoder::start(const char* server, const char* id, int channel)
{
    strcpy(mID, id);
    mRtmpReader->open(server, id, channel);
    mChannel = channel;
    mNeedKeyFrame = true;
    mState = STAT_IDLE;
}

void CVideoDecoder::stop()
{
    mState = STAT_STOP;
}

void CVideoDecoder::swapChannel(int channel)
{
    if (mRtmpReader->isConnected() && channel != mChannel) {
        char stream[128];
        sprintf(stream, "%s%d", mID, channel);
        int ret = mRtmpReader->swapStream(stream);
        LOGI("swapStream('%s') = %d", stream, ret);
        mNeedKeyFrame = true;
        mChannel = channel;
    }
}

void CVideoDecoder::saveStream(const char* serial_no)
{
    char stream_name[64];
    sprintf(stream_name, "%s0", mID);
    mRtmpReader->saveStream(stream_name, serial_no);
}

void CVideoDecoder::registerDecoder(create_decoder_cb create_cb, destroy_decoder_cb destroy_cb, decode_frame_cb decode_cb)
{
    mOnDecoderCreateHandler = create_cb;
    mOnDecoderDestroyHandler = destroy_cb;
    mOnDecoderDecodeHandler = decode_cb;
}

char* CVideoDecoder::getDecodedFrame()
{
    char* frame = NULL;
    if (mFrameCount > 0) {
        MUTEX_LOCK(mLock, {
            frame = mFrames[mFrameRindex];
            mFrames[mFrameRindex] = NULL;
            mFrameRindex = (mFrameRindex + 1) % MAX_FRAMES;
            mFrameCount--;
        });
    }
    return frame;
}

void CVideoDecoder::freeDecodedFrame(char* frameptr)
{
    safe_freep(frameptr);
    mCanNotifyRenderer = true;
}

void CVideoDecoder::onThread()
{
    mIsRunning = true;
    mCanNotifyRenderer = true;
    
    avcodec_register_all();
    av_log_set_level(AV_LOG_WARNING);
    av_log_set_callback(log_callback_report);
    
    memset(&mStatisInfo, 0, sizeof(mStatisInfo));
    while (mIsRunning) {
        bool can_sleep = false;
        
        switch(mState) {
            case STAT_IDLE:
            case STAT_ERROR:
                usleep(1000 * 50);
                break;
            case STAT_START:
                mRtmpReader->open(mServer, mID, mChannel);
                mState = STAT_IDLE;
                break;
            case STAT_STOP:
                mRtmpReader->close();
                destroyDecoder();
                mState = STAT_IDLE;
                break;
            case STAT_READY:
                if (!decodeFrame()) {
                    can_sleep = true;
                }
                break;
        };
        
        if (mState == STAT_READY && mFrameCount > 0 && mCanNotifyRenderer) {
            mCanNotifyRenderer = false;
            if (mOnFrameAvailableHandler)
                mOnFrameAvailableHandler();
        }
        else if (can_sleep) {
            usleep(1000 * 10);
        }
        
        if (mState == STAT_READY) {
            int64_t _now = sys_gettime();
            if (mStatisInfo.last_time == 0) {
                mRtmpReader->getReceivedInfo(&mStatisInfo.rcvbytes, &mStatisInfo.rcvpkts);
                mStatisInfo.last_time = _now;
            }
            else if (_now - mStatisInfo.last_time >= 1000) {
                mRtmpReader->getReceivedInfo(&mStatisInfo.rcvbytes, &mStatisInfo.rcvpkts);
                
                int bitrate = (mStatisInfo.rcvbytes * 8) / 1000;
                if (mStatisInfo.max_bitrate < bitrate)
                    mStatisInfo.max_bitrate = bitrate;
                if (mOnStatisInfoHandler)
                    mOnStatisInfoHandler(bitrate, mStatisInfo.max_bitrate, mStatisInfo.rcvpkts, mStatisInfo.dec_pkts, mStatisInfo.discard_pkts);
                LOGI("[STATIS] bitrate: %d kbps, max_bitrate: %d kbps, recv %d pkts, decoded %d pkts, discard %d pkts",
                     bitrate, mStatisInfo.max_bitrate, mStatisInfo.rcvpkts, mStatisInfo.dec_pkts, mStatisInfo.discard_pkts);
                
                mStatisInfo.dec_pkts = 0;
                mStatisInfo.discard_pkts = 0;
                mStatisInfo.last_time = _now;
            }
        }
    }
    
    delete mRtmpReader;
    mRtmpReader = NULL;
    
    destroyDecoder();
    freeAllPackets();
    freeAllFrames();
    mThread = 0;
}

void CVideoDecoder::onVideoPacket(char* data, int size)
{
    //LOGI("-->onVideoPacket -- %p, %d", data, size);
    // 获取PPS && SPS
    if (data[0] == 0x17 && data[1] == 0x00) {
        uint8_t extradata[128], *p = extradata, *p1;
        int extradata_size, len;
        
        *p++ = 0; *p++ = 0; *p++ = 0; *p++ = 1;
        len = ((((int)data[11]) << 8) | ((int)data[12]));
        memcpy(p, data + 13, len);        // 复制SPS
        p += len;
        
        *p++ = 0; *p++ = 0; *p++ = 0; *p++ = 1;
        
        p1 = (uint8_t*)data + 13 + len + 1;
        len = ((((int)*p1) << 8) | (int)*(p1 + 1));
        memcpy(p, p1 + 2, len);        // 复制PPS
        
        p += len;
        extradata_size = p - extradata;
        
        if (mExtradataSize != extradata_size || memcmp(mExtradata, extradata, extradata_size) != 0) {
            // 视频SPS、PPS数据不同，需要重新创建解码器
            destroyDecoder();
            if (!createDecoder(extradata, extradata_size, data, size)) {
                mRtmpReader->close();
                mState = STAT_IDLE;
                LOGE("FAILED to create decoder!!!");
            }
            else {
                mState = STAT_READY;
                mCanNotifyRenderer = true;
                memcpy(mExtradata, extradata, extradata_size);
                mExtradataSize = extradata_size;
                
                if (mOnVideoReadyHandler != NULL)
                    mOnVideoReadyHandler(mVideoWidth, mVideoHeight);
            }
        }
        payload_freep(data);
        return;
    }
    
    // 保存视频帧
    bool key = (data[0] == 0x17);
    
    if ((mDecodeCtx == NULL && !mUseExtraDecoder) || (mNeedKeyFrame && !key)) {
        payload_freep(data);
        LOGI("skip video packet! (%p) (key %d)", mDecodeCtx, key);
        return;
    }
    
    MUTEX_LOCK(mLock, {
        if (key) {
            // 收到关键帧，将之前未解码的帧都删除
            for (int i = 0; i < mPacketCount; i++) {
                payload_freep(mPackets[(i + mPacketRindex) % MAX_PACKETS].data);
            }
            memset(mPackets, 0, sizeof(mPackets));
            mPacketRindex = 0;
            mPacketWindex = 0;
            mPacketCount = 0;
        }
        
        if (mPacketCount >= MAX_PACKETS) {
            LOGE("*** Packet Buffer full, skip packet!!! ***");
        }
        else {
            mPackets[mPacketWindex].data = data;
            mPackets[mPacketWindex].size = size;
            mPackets[mPacketWindex].key = key;
            mPacketWindex = (mPacketWindex + 1) % MAX_PACKETS;
            mPacketCount++;
        }
    });
    
    mNeedKeyFrame = false;
}

bool CVideoDecoder::createDecoder(uint8_t* extradata, int extradata_size, char* payload, int payload_size)
{
    mUseExtraDecoder = false;
    mUseMediaCodec = false;
    mPts = 0;
    
    if (srs_parse_videosize_from_sps(payload, payload_size, &mVideoWidth, &mVideoHeight) == 0) {
        LOGI("*** Parse Video Size from SPS is %d x %d ***", mVideoWidth, mVideoHeight);
    }
    
    if (mOnDecoderCreateHandler != NULL) {
        // 调用外部解码器接口创建解码器
        if (mOnDecoderCreateHandler(extradata, extradata_size, mVideoWidth, mVideoHeight)) {
            mUseExtraDecoder = true;
            LOGI("CREATE extra decoder success!!!");
            return true;
        }
    }
    
    AVCodecContext* ctx;
    AVCodec* codec = NULL;
    int use_hard_decoder = 1;
    int ret;
    
#if defined(ANDROID) && 0
    codec = avcodec_find_decoder_by_name("h264_mediacodec");
#endif
    
    do {
        if (codec == NULL) {
            use_hard_decoder = 0;
            codec = avcodec_find_decoder(AV_CODEC_ID_H264);
        }
        
        ctx = avcodec_alloc_context3(codec);
        if (!ctx) {
            return false;
        }
        
        ctx->extradata = (uint8_t*)av_malloc(extradata_size);
        memcpy(ctx->extradata, extradata, extradata_size);
        ctx->extradata_size = extradata_size;
        ctx->profile = FF_PROFILE_H264_BASELINE;
        ctx->flags |= AV_CODEC_FLAG_LOW_DELAY;
        ctx->flags2 |= AV_CODEC_FLAG2_FAST | CODEC_FLAG2_CHUNKS;
        ctx->pix_fmt = AV_PIX_FMT_YUV420P;
        
        if ((ret = avcodec_open2(ctx, codec, NULL)) < 0) {
            LOGE("Failed to open ffmpeg decoder!!!! error: %d", ret);
            avcodec_free_context(&ctx);
            if (use_hard_decoder) {
                ctx = NULL;
                codec = NULL;
                continue;
            }
            return false;
        }
        
#ifdef ANDROID
        if (ctx->pix_fmt != AV_PIX_FMT_YUV420P) {
            mSwsCtx = sws_getContext(mVideoWidth, mVideoHeight, ctx->pix_fmt,
                                     mVideoWidth, mVideoHeight, AV_PIX_FMT_YUV420P,
                                     SWS_BILINEAR, NULL, NULL, NULL);
            if (mSwsCtx == NULL) {
                LOGE("Failed to create pixel format convert!!!");
                avcodec_free_context(&ctx);
                if (use_hard_decoder) {
                    ctx = NULL;
                    codec = NULL;
                    continue;
                }
                return false;
            }
        }
#endif
        
        if (use_hard_decoder) {
            LOGI("^_^ HARD DECODER is using!!! ^_^");
            mUseMediaCodec = true;
            mFirstInputFrame = true;
        }
        LOGI("FFMPEG DECODER CREATED.");
    } while (0);
    
    mDecodeCtx = ctx;
    return true;
}

void CVideoDecoder::destroyDecoder()
{
    if (mUseExtraDecoder) {
        mOnDecoderDestroyHandler();
        mUseExtraDecoder = false;
    }
    else {
        if (mDecodeCtx != NULL) {
            avcodec_free_context(&mDecodeCtx);
            mDecodeCtx = NULL;
        }
#ifdef ANDROID
        if (mSwsCtx != NULL) {
            sws_freeContext(mSwsCtx);
            mSwsCtx = NULL;
        }
#endif
    }
    mExtradataSize = 0;
    
    freeAllPackets();
    freeAllFrames();
}

bool CVideoDecoder::decodeFrame()
{
    struct video_packet packets[MAX_PACKETS];
    int nb_packets = 0;
    
    if (mUseExtraDecoder) {
        // 调用外部解码器刷新解码帧
        char* frame = NULL;
        if (mOnDecoderDecodeHandler(NULL, 0, &frame) == 1) {
            if (frame != NULL) {
                putDecodedFrame(frame);
            }
            return true;
        }
    }
    else if (mUseMediaCodec) {
        AVFrame* frame = av_frame_alloc();
        if (avcodec_receive_frame(mDecodeCtx, frame) == 0) {
            putDecodedAVFrame(frame);
            return true;
        }
        else
            av_frame_free(&frame);
    }
    
    if (mPacketCount == 0)
        return false;
    
    MUTEX_LOCK(mLock, {
        if (mPacketCount <= 4)
            nb_packets = 1;
        else if (mPacketCount >= 10)
            nb_packets = mPacketCount - 2;
        else
            nb_packets = 2;
        
        int widx = 0;
        for (int i = 0; i < nb_packets; i++) {
            if (mPackets[mPacketRindex].data != NULL) {
                memcpy(&packets[widx++], &mPackets[mPacketRindex], sizeof(struct video_packet));
                mPackets[mPacketRindex].data = NULL;
            }
            mPacketRindex = (mPacketRindex + 1) % MAX_PACKETS;
            mPacketCount--;
        }
        nb_packets = widx;
    });
    if (nb_packets == 0)
        return true;
    
    if (mUseExtraDecoder) {
        // 使用外部解码器解码
        char* frame = NULL;
        for (int i = 0; i < nb_packets; i++) {
            char* new_frame = NULL;
            uint8_t* data = (uint8_t*)packets[i].data + 5;
            data[0] = 0; data[1] = 0; data[2] = 0; data[3] = 1;
            int ret = mOnDecoderDecodeHandler(data, packets[i].size - 5, &new_frame);
            payload_freep(packets[i].data);
            if (ret == 1 && new_frame != NULL) {
                safe_freep(frame);
                frame = new_frame;
            }
            mStatisInfo.dec_pkts++;
        }
        if (frame != NULL) {
            putDecodedFrame(frame);
        }
    }
    else {
        // 使用ffmpeg解码
        AVFrame* frame = av_frame_alloc();
        if (mUseMediaCodec && mFirstInputFrame && nb_packets < 4) {
            for (int i = 0; i < 5; i++) {
                AVPacket pkt;
                memset(&pkt, 0, sizeof(pkt));
                pkt.data = (uint8_t*)packets[0].data + 5;
                pkt.size = packets[0].size - 5;
                pkt.pts = mPts++;
                pkt.data[0] = 0;
                pkt.data[1] = 0;
                pkt.data[2] = 0;
                pkt.data[3] = 1;
                avcodec_send_packet(mDecodeCtx, &pkt);
                int ret = avcodec_receive_frame(mDecodeCtx, frame);
                if (ret == 0) {
                    break;
                }
            }
            av_frame_free(&frame);
        }
        av_frame_free(&frame);
        for (int i = 0; i < nb_packets; i++) {
            AVPacket pkt;
            memset(&pkt, 0, sizeof(pkt));
            pkt.data = (uint8_t*)packets[i].data + 5;
            pkt.size = packets[i].size - 5;
            pkt.pts = mPts++;
            pkt.data[0] = 0;
            pkt.data[1] = 0;
            pkt.data[2] = 0;
            pkt.data[3] = 1;
            avcodec_send_packet(mDecodeCtx, &pkt);
            payload_freep(packets[i].data);
            
            AVFrame* new_frame = av_frame_alloc();
            if (new_frame == NULL) {
                i++;
                while (i < nb_packets) {
                    payload_freep(packets[i].data);
                    i++;
                }
                break;
            }
            int ret = avcodec_receive_frame(mDecodeCtx, new_frame);
            if (ret == 0) {
                av_frame_free(&frame);
                frame = new_frame;
            }
            else
                av_frame_free(&new_frame);
        }
        
        if (frame != NULL) {
            mFirstInputFrame = false;
            putDecodedAVFrame(frame);
            mStatisInfo.dec_pkts++;
            mStatisInfo.discard_pkts += nb_packets - 1;
        }
        else
            mStatisInfo.discard_pkts += nb_packets;
    }
    
    return true;
}

void CVideoDecoder::putDecodedAVFrame(AVFrame* frame)
{
    char* framedata = (char*)malloc(mVideoWidth * mVideoHeight * 3 / 2);
    if (framedata != NULL) {
        int size1 = mVideoWidth * mVideoHeight;
        int size2 = size1 / 4;
#ifdef ANDROID
        if (mUseMediaCodec && mSwsCtx != NULL) {
            uint8_t* data[4];
            int linesize[4];
            linesize[0] = mVideoWidth;
            linesize[1] = linesize[2] = mVideoWidth / 2;
            linesize[3] = 0;
            
            data[0] = (uint8_t*)framedata;
            data[1] = (uint8_t*)framedata + size1;
            data[2] = (uint8_t*)framedata + (size1 + size2);
            data[3] = 0;
            sws_scale(mSwsCtx, frame->data, frame->linesize,
                      0, mVideoHeight, data, linesize);
        }
        else
#endif
        {
            memcpy(framedata, frame->data[0], size1);
            memcpy(framedata + size1, frame->data[1], size2);
            memcpy(framedata + size1 + size2, frame->data[2], size2);
        }
        putDecodedFrame(framedata);
    }
    else {
        LOGE("FAILED to alloc decoded frame data!!!");
    }
    av_frame_free(&frame);
}

void CVideoDecoder::putDecodedFrame(char* frameptr)
{
    MUTEX_LOCK(mLock, {
        if (mFrameCount >= MAX_FRAMES) {
            safe_freep(mFrames[mFrameRindex]);
            mFrameRindex = (mFrameRindex + 1) % MAX_FRAMES;
            mFrameCount--;
        }
        mFrames[mFrameWindex] = frameptr;
        mFrameWindex = (mFrameWindex + 1) % MAX_FRAMES;
        mFrameCount++;
    });
}

void CVideoDecoder::freeAllFrames()
{
    MUTEX_LOCK(mLock, {
        for (int i = 0; i < mFrameCount; i++) {
            safe_freep(mFrames[(i + mFrameRindex) % MAX_FRAMES]);
        }
        mFrameRindex = 0;
        mFrameWindex = 0;
        mFrameCount = 0;
    });
}

void CVideoDecoder::freeAllPackets()
{
    MUTEX_LOCK(mLock, {
        for (int i = 0; i < mPacketCount; i++) {
            payload_freep(mPackets[(i + mPacketRindex) % MAX_PACKETS].data);
        }
        memset(mPackets, 0, sizeof(mPackets));
        mPacketRindex = 0;
        mPacketWindex = 0;
        mPacketCount = 0;
    });
}

#ifdef __cplusplus
}
#endif
